#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：Flower_Class 
@File    ：detect.py
@IDE     ：PyCharm 
@Author  ：Rice_cc
@Date    ：2024/10/16 2:31 
'''

import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import tensorflow as tf


#加载数据，按照8:2的比例加载花卉数据
def data_load(data_dir, img_height, img_width, batch_size):
    train_ds = tf.keras.preprocessing.image_dataset_from_directory(data_dir,
                                                                   label_mode='categorical',
                                                                   validation_split=0.2,
                                                                   subset="training",
                                                                   seed=123,
                                                                   image_size=(img_height, img_width),
                                                                   batch_size=batch_size)
    val_ds = tf.keras.preprocessing.image_dataset_from_directory(
        data_dir,
        label_mode='categorical',
        validation_split=0.2,
        subset="validation",
        seed=123,
        image_size=(img_height, img_width),
        batch_size=batch_size)

    class_names = train_ds.class_names
    return train_ds, val_ds, class_names


def test(is_transfer=True):
    train_ds, val_ds, class_names = data_load("C:/Users/23987/Downloads/flower_photos/flower_split/test", 224, 224, 4)
    model = tf.keras.models.load_model("models/mobilenet_flower.h5")
    model.summary()
    loss, accuracy = model.evaluate(val_ds)
    print('Test accuracy :', accuracy)


if __name__ == '__main__':
    test(True)
